package denoise;

import java.awt.*;
import java.awt.image.*;
import java.io.*;
import java.util.*;

import javax.swing.*;
import javax.imageio.*;

import denoise.util.*;

public class DenoiseGrader implements ParameterListener
{
    JFrame jf;
    JImage jimLeft, jimRight;

    ParameterGUI pg;

    DenoiseAgent agent;
    ArrayList<TestCase> testcases;
    int testidx = 0;

    public static class TestCase
    {
        BufferedImage original;
        BufferedImage noisy;
	BufferedImage[] training;
        double noise;

        public TestCase(BufferedImage original, double noise, BufferedImage[] training_data)
        {
            this.original = original;
            this.noise = noise;

            int width = original.getWidth(), height = original.getHeight();
            noisy = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);

            Random r = new Random(0);

            for (int y = 0; y < height; y++) {
                for (int x = 0; x < width ;x++) {

                    if (r.nextDouble() < noise)
                        noisy.setRGB(x, y, original.getRGB(x, y)^0xffffff);
                    else
                        noisy.setRGB(x, y, original.getRGB(x, y));
                }
            }
	    training = training_data;
        }
    }

    public DenoiseGrader(ArrayList<TestCase> testcases, DenoiseAgent agent, boolean gui)
    {
        this.testcases = testcases;
        this.agent = agent;

        if (!gui) {
            new RunThread(Integer.MAX_VALUE).start();
            return;
        }

        pg = new ParameterGUI();
        pg.addButtons("next", "next");
        pg.addButtons("all", "all");

        jf = new JFrame("Denoise Grader");

        jimLeft = new JImage();
        jimRight = new JImage();

        jf.setLayout(new BorderLayout());
        JSplitPane jsp = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, jimLeft, jimRight);
        jsp.setDividerLocation(.5);
        jsp.setResizeWeight(.5);
        jf.add(jsp, BorderLayout.CENTER);
        jf.add(pg, BorderLayout.SOUTH);

        jf.setSize(800,600);
        jf.setVisible(true);

        pg.addListener(this);

        step();
    }

    double mseM; // sum of MSE
    int    mseN; // number of images.
    double totalTime;

    synchronized void step()
    {
        if (testidx == testcases.size()) {
            // print summary and exit.
            System.out.printf("SUMMARY : total time   %-15.5f s, MSE %-15.5f\n", totalTime, mseM/ mseN);
            System.exit(0);
        }

        TestCase tc = testcases.get(testidx);

        if (jimLeft != null)
            jimLeft.setImage(tc.noisy);

        long starttime = System.currentTimeMillis();
        BufferedImage denoised = agent.denoise(tc.noisy, tc.training);
        long endtime = System.currentTimeMillis();

        double dt = (endtime - starttime) / 1000.0;
        totalTime += dt;

        if (jimRight != null)
            jimRight.setImage(denoised);

        // compute score.
        int width = denoised.getWidth(), height = denoised.getHeight();

        double mse = 0;
        double N = 0;
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width ;x++) {

                int v = tc.original.getRGB(x, y) & 0x0000ff;
                int w = denoised.getRGB(x, y) & 0x0000ff;

                mse += (v-w)*(v-w);
                N++;
            }
        }

        mse /= N;

        System.out.printf("image %2d: elapsed time %-15.5f s, MSE %-15.5f\n", testidx, dt, mse);
        mseM += mse;
        mseN ++;

        testidx++;
    }

    class RunThread extends Thread
    {
        int nsteps;

        RunThread(int nsteps)
        {
            this.nsteps = nsteps;
        }

        public void run()
        {
            for (int i = 0; i < nsteps; i++)
                step();
        }
    }

    public void parameterChanged(ParameterGUI pg, String name)
    {
        if (name.equals("next"))
            new RunThread(1).start();
        if (name.equals("all"))
            new RunThread(Integer.MAX_VALUE).start();
    }

    public static BufferedImage binarize(BufferedImage in)
    {
        int width = in.getWidth(), height = in.getHeight();
        BufferedImage out = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);

        Random r = new Random(0);

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width ;x++) {

                int v = in.getRGB(x, y) & 0x00ff00; // pull out green channel.
                if (v >= 0x8000)
                    out.setRGB(x,y, 0xffffff);
                else
                    out.setRGB(x,y, 0x000000);
            }
        }

        return out;
    }

    public static void main(String args[])
    {
        ArrayList<TestCase> testcases = new ArrayList<TestCase>();
        GetOpt gopt = new GetOpt();
        gopt.addString('t', "tests", "testcases.txt", "File containing testcase descriptions");
        gopt.addString('a', "agent", "denoise.ExampleAgent", "Classpath of DenoiseAgent");
        gopt.addBoolean('g', "gui",  true, "Enable GUI");
        gopt.addBoolean('h', "help", false, "Show help");

        gopt.parse(args);

        if (gopt.getBoolean("help")) {
            gopt.doHelp();
            System.exit(0);
        }

        String testpath = gopt.getString("tests");

        DenoiseAgent agent = (DenoiseAgent) ReflectUtil.createObject(gopt.getString("agent"));

        try {
            BufferedReader ins = new BufferedReader(new FileReader(testpath));
            String line;

            while ((line = ins.readLine())!=null) {

                line = line.trim();
                if (line.startsWith("#"))
                    continue;
                String toks[] = line.split("\\s+");
                if (toks.length != 4) {
                    System.out.println("Illegal line: "+line);
                    continue;
                }

                double noise = Double.parseDouble(toks[1]);
		System.out.printf("Reading '%s', will add noise %f\n", toks[0], noise);
	
		System.out.println("going to read "+toks[2] + " which contains "+ toks[3] + " lines");
		int num_inputs = Integer.parseInt(toks[3]);
		BufferedImage[] training_data = new BufferedImage[num_inputs];

                
                BufferedImage im = ImageIO.read(new File(toks[0]));
                im = binarize(im);

		BufferedReader training_br = new BufferedReader(new FileReader(toks[2]));
		String input;
		int ctr = 0;
		while((input = training_br.readLine())!= null && ctr++ < num_inputs)
		{
			input = input.trim();
			BufferedImage im1 = ImageIO.read(new File(input));
			im1 = binarize(im1);
			training_data[ctr-1] = im1;
		}

                TestCase tc = new TestCase(im, noise, training_data);
                testcases.add(tc);
            }

        } catch (IOException ex) {
            System.out.println("ex: "+ex);
        }

        new DenoiseGrader(testcases, agent, gopt.getBoolean("gui"));
    }
}
